package lock

import (
	"bytes"
	"runtime"
	"strconv"
	"sync"
	"sync/atomic"
)

const (
	Goroutine = "goroutine" // 协程字符串
	Space     = " "         // 空格
	CharSpace = ' '         // 空格
)

// GetCurrentGid 获取当前协程id
func GetCurrentGid() int64 {
	bs := make([]byte, 64)
	bs = bs[:runtime.Stack(bs, false)]
	bs = bytes.TrimPrefix(bs, []byte(Goroutine+Space))
	bs = bs[:bytes.IndexByte(bs, CharSpace)]
	gid, err := strconv.Atoi(string(bs))
	if err != nil {
		panic(err)
	}
	return int64(gid)
}

// ReentrantMutex 实现可重入锁
type ReentrantMutex struct {
	mutex          sync.Mutex // 互斥锁
	gid            int64      // 协程id
	reentrantCount int32      // 同一个协程的重入次数
}

func (my *ReentrantMutex) Lock() {
	gid := GetCurrentGid()
	if atomic.LoadInt64(&my.gid) == gid {
		atomic.AddInt32(&my.reentrantCount, 1)
		return
	}
	my.mutex.Lock()
	atomic.StoreInt64(&my.gid, gid)
	my.reentrantCount = 1
}

func (my *ReentrantMutex) Unlock() {
	gid := GetCurrentGid()
	if atomic.LoadInt64(&my.gid) != gid {
		panic("当前协程不存在需要解锁的锁")
	}

	atomic.AddInt32(&my.reentrantCount, -1)
	if atomic.LoadInt32(&my.reentrantCount) != 0 {
		return
	}

	atomic.StoreInt64(&my.gid, -1)
	my.mutex.Unlock()
}

func (my *ReentrantMutex) TryLock() bool {
	gid := GetCurrentGid()
	if atomic.LoadInt64(&my.gid) == gid {
		atomic.AddInt32(&my.reentrantCount, 1)
		return true
	}

	if my.mutex.TryLock() {
		atomic.StoreInt64(&my.gid, gid)
		my.reentrantCount = 1
		return true
	}
	return false
}
